#include "Mem2Reg.hpp"
#include "IRBuilder.hpp"
#include "Value.hpp"
#include "logging.hpp"
#include <memory>
#include "CodeGenUtil.hpp"
#include "lightir/Value.hpp"
#include "lightir/Instruction.hpp"

/**
 * @brief Mem2Reg Pass的主入口函数
 * 
 * 该函数执行内存到寄存器的提升过程，将栈上的局部变量提升到SSA格式。
 * 主要步骤：
 * 1. 创建并运行支配树分析
 * 2. 对每个非声明函数：
 *    - 清空相关数据结构
 *    - 插入必要的phi指令
 *    - 执行变量重命名
 * 
 * 注意：函数执行后，冗余的局部变量分配指令将由后续的死代码删除Pass处理
 */
void Mem2Reg::run()
{
    // 创建支配树分析 Pass 的实例
    dominators_ = std::make_unique<Dominators>(m_);
    // 建立支配树
    dominators_->run();
    // 以函数为单元遍历实现 Mem2Reg 算法
    for (auto &f : m_->get_functions())
    {
        if (f.is_declaration())
            continue;
        func_ = &f;
        if (func_->get_basic_blocks().size() >= 1)
        {
            // 预处理alloc指令，适配可达定义分析
            loop_alloc_inv_hoist();
            // 对应伪代码中 phi 指令插入的阶段
            generate_phi();
            // 对应伪代码中重命名阶段
            rename(func_->get_entry_block());
        }
        // 后续 DeadCode 将移除冗余的局部变量的分配空间
    }
}

/**
 * @brief 将所有基本块中的 alloca 指令提取出来，统一移动到函数的入口基本块的起始位置
 */
void Mem2Reg::loop_alloc_inv_hoist()
{
    std::vector<Instruction*> alloc_vec;
    for(auto &bb : func_->get_basic_blocks()){
        auto b = &bb;
        std::vector<Instruction*> replace_inst;
        for(auto &inst : b->get_instructions()){
            auto instr = &inst;
            if(instr->is_alloca()){
                alloc_vec.push_back(instr);
                replace_inst.push_back(instr);
            }
        }
        for(auto instr : replace_inst){
            instr->get_parent()->remove_instr(instr);
        }
    }
    auto entry_bb = func_->get_entry_block();
    for(auto &instr : alloc_vec){
        entry_bb->add_instr_begin(instr);
    }
}

/**
 * @brief 在必要的位置插入phi指令
 * 
 * 该函数实现了经典的phi节点插入算法：
 * 1. 收集全局活跃变量：
 *    - 扫描所有store指令
 *    - 识别在多个基本块中被赋值的变量
 * 
 * 2. 插入phi指令：
 *    - 对每个全局活跃变量
 *    - 在其定值点的支配边界处插入phi指令
 *    - 使用工作表法处理迭代式的phi插入
 * 
 * phi指令的插入遵循最小化原则，只在必要的位置插入phi节点
 */
void Mem2Reg::generate_phi()
{
    // 步骤一：找到活跃在多个 block 的全局名字集合
    std::set<Value*> alloca_var; // 存储活跃于多个基本块中的alloca变量
    std::map<Value*, std::set<BasicBlock*>> alloca_var_map; // 记录变量使用的基本块
    
    // 扫描所有store指令，收集活跃变量
    for (auto& bb : func_->get_basic_blocks()) {
        auto b = &bb;
        for (auto& inst : b->get_instructions()) {
            if (inst.is_store()) {
                auto store_inst = static_cast<StoreInst*>(&inst);
                auto lval = store_inst->get_lval();
                // 只处理有效的alloca指针（非GEP、非全局变量）
                if (is_valid_ptr(lval)) {
                    alloca_var.insert(lval);
                    alloca_var_map[lval].insert(b);
                }
            }
        }
    }
    
    // 步骤二：在支配边界插入phi指令
    for (auto alloca : alloca_var) {
        std::set<BasicBlock*> work_list = alloca_var_map[alloca];
        std::set<BasicBlock*> has_phi; // 记录已经插入phi的基本块
        
        while (!work_list.empty()) {
            // 从工作表取出一个基本块
            BasicBlock* bb = *work_list.begin();
            work_list.erase(work_list.begin());
            
            // 获取该基本块的支配边界
            const auto& df = dominators_->get_dominance_frontier(bb);
            
            for (auto d : df) {
                // 如果支配边界中的基本块还没有该变量的phi指令
                if (has_phi.find(d) == has_phi.end()) {
                    // 创建phi指令
                    auto phi_type = alloca->get_type()->get_pointer_element_type();
                    auto phi = PhiInst::create_phi(phi_type, d);
                    phi->set_lval(alloca); // 设置关联的alloca变量
                    
                    // 添加到基本块开头
                    d->add_instr_begin(phi);
                    has_phi.insert(d);
                    
                    // 记录这个phi指令对应的alloca变量
                    phi_instructions_[alloca].insert(phi);
                    
                    // 如果该基本块之前没有store过这个变量，加入工作表
                    if (alloca_var_map[alloca].find(d) == alloca_var_map[alloca].end()) {
                        work_list.insert(d);
                    }
                }
            }
        }
    }
}

void Mem2Reg::rename(BasicBlock *bb)
{
    std::vector<Instruction *> deleted_instructions;
    
    // 步骤三：将 phi 指令作为 lval 的最新定值
    for (auto& instr : bb->get_instructions()) {
        if (instr.is_phi()) {
            auto phi = dynamic_cast<PhiInst*>(&instr);
            Value* alloca = phi->get_lval();
            var_stack[alloca].push_back(phi);
        }
    }
    
    // 步骤四：用 lval 最新的定值替代对应的load指令
    for (auto& instr : bb->get_instructions()) {
        if (instr.is_load() && is_valid_ptr(instr.get_operand(0))) {
            Value* alloca = instr.get_operand(0);
            if (!var_stack[alloca].empty()) {
                Value* replacement = var_stack[alloca].back();
                instr.replace_all_use_with(replacement);
                deleted_instructions.push_back(&instr);
            }
        }
    }
    
    // 步骤五：将 store 指令的 rval 作为 lval 的最新定值
    for (auto& instr : bb->get_instructions()) {
        if (instr.is_store() && is_valid_ptr(instr.get_operand(1))) {
            Value* alloca = instr.get_operand(1);
            Value* value = instr.get_operand(0);
            var_stack[alloca].push_back(value);
            deleted_instructions.push_back(&instr);
        }
    }
    
    // 步骤六：为 lval 对应的 phi 指令参数补充完整，处理分支条件
    for (auto succ_bb : bb->get_succ_basic_blocks()) {
        for (auto& instr : succ_bb->get_instructions()) {
            if (instr.is_phi()) {
                auto phi = dynamic_cast<PhiInst*>(&instr);
                Value* alloca = phi->get_lval();
                if (!var_stack[alloca].empty()) {
                    Value* current_val = var_stack[alloca].back();
                    phi->add_operand(current_val, bb); // 传递当前基本块作为来源
                } else {
                    // 如果没有定义，添加一个undef值
                    auto undef = UndefValue::get(phi->get_type());
                    phi->add_operand(undef, bb);
                }
            }
            
            // 处理分支指令中的条件操作数（解决分支条件变量未提升问题）
            if (instr.is_br() && instr.get_num_operand() == 3) {
                Value* cond = instr.get_operand(0);
                if (is_valid_ptr(cond)) {
                    Value* alloca = cond;
                    if (!var_stack[alloca].empty()) {
                        Value* replacement = var_stack[alloca].back();
                        
                        // 使用set_operand替换第0个操作数（替代不存在的replace_operand）
                        instr.set_operand(0, replacement);
                    }
                }
            }
        }
    }
    
    // 步骤七：对 bb 在支配树上的所有后继节点，递归执行 rename 操作
    const auto& children = dominators_->get_dom_tree_succ_blocks(bb);
    for (auto child : children) {
        rename(child);
    }
    
    // 步骤八：pop出 lval 的最新定值
    for (auto& instr : bb->get_instructions()) {
        if (instr.is_phi()) {
            auto phi = dynamic_cast<PhiInst*>(&instr);
            Value* alloca = phi->get_lval();
            if (!var_stack[alloca].empty() && var_stack[alloca].back() == phi) {
                var_stack[alloca].pop_back();
            }
        }
        if (instr.is_store() && is_valid_ptr(instr.get_operand(1))) {
            Value* alloca = instr.get_operand(1);
            if (!var_stack[alloca].empty()) {
                var_stack[alloca].pop_back();
            }
        }
    }
    
    // 步骤九：清除冗余的指令
    for (auto inst : deleted_instructions)
    {
        bb->erase_instr(inst);
    }
}

